package com.ktjiaoyu.travelaiagent.app;

import com.ktjiaoyu.travelaiagent.advisors.TokenCounterAdvisor;
import com.ktjiaoyu.travelaiagent.mcp.McpToolCallbackProxy;
import com.ktjiaoyu.travelaiagent.memory.KryoChatMemory;

import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.MessageChatMemoryAdvisor;
import org.springframework.ai.chat.client.advisor.QuestionAnswerAdvisor;
import org.springframework.ai.chat.client.advisor.api.Advisor;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.chat.prompt.SystemPromptTemplate;
import org.springframework.ai.model.function.FunctionCallback;
import org.springframework.ai.tool.ToolCallback;
import org.springframework.ai.tool.ToolCallbackProvider;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.core.io.Resource;
import org.springframework.stereotype.Component;
import reactor.core.publisher.Flux;

import java.util.Map;

import static org.springframework.ai.chat.client.advisor.AbstractChatMemoryAdvisor.CHAT_MEMORY_CONVERSATION_ID_KEY;
import static org.springframework.ai.chat.client.advisor.AbstractChatMemoryAdvisor.CHAT_MEMORY_RETRIEVE_SIZE_KEY;

@Component
@Slf4j
public class TravelApp {

    // 定义ChatClient对象为final对象，则必须在构造函数中初始化
    private final ChatClient chatClient;
    // 在构造函数中注入dashscopeChatModel
    public TravelApp(ChatModel dashscopeChatModel,
                     // 由@Qualifier注解指定使用的 RAG Advisor
                     @Qualifier("cloudRagAdvisor") Advisor ragAdvisor,
                     @Value("classpath:/prompts/travel-app-system.pt") Resource systemPromptResource,
                     @Value("${travel.app.name:郴州科泰}") String name) {
        SystemPromptTemplate systemPromptTemplate = new SystemPromptTemplate(systemPromptResource);
        Map<String, Object> vars = Map.of("name", name);
        Prompt systemPrompt = systemPromptTemplate.create(vars);
        String systemPromptStr = systemPrompt.getContents();
        // ChatMemory chatMemory = new InMemoryChatMemory();
        // 修改为Kryo持久化存储的聊天记忆策略，savePath可以通过配置文件配置
        String savePath = System.getProperty("user.dir") + "/chat-memory";
        ChatMemory chatMemory = new KryoChatMemory(savePath);
        this.chatClient = ChatClient.builder(dashscopeChatModel)
                .defaultAdvisors(
                        // 添加知识库问答顾问
                        ragAdvisor,
                        // 定义消息记忆策略
                        new MessageChatMemoryAdvisor(chatMemory),
                        // 创建并注册自定义的Token计数顾问
                        new TokenCounterAdvisor()
                )
                .defaultSystem(systemPromptStr)
                .build();
    }

    // 注入工具数组（注意Resource类和注解有冲突，此处需要写全名）
    @jakarta.annotation.Resource
    private ToolCallback[] aiTools;
    // 配置的MCP服务会被自动解析
    @jakarta.annotation.Resource
    private ToolCallbackProvider toolCallbackProvider;

    /**
     * 与AI大模型对话接口
     * @param chatId 对话ID
     * @param message 消息内容
     * @return 大模型响应结果
     */
    public String chat(String chatId, String message) {
        // 从toolCallbackProvider获得MCP服务产生的ToolCallback数组
        FunctionCallback[] mcpToolCallbacks = toolCallbackProvider.getToolCallbacks();
        ChatResponse chatResponse = chatClient.prompt()
                .user(message)
                .advisors(spec -> spec
                        .param(CHAT_MEMORY_CONVERSATION_ID_KEY, chatId)
                        .param(CHAT_MEMORY_RETRIEVE_SIZE_KEY, 10)
                )
                // 注册可以使用的工具
                .tools(aiTools)
                // 注册代理后的MCP服务工具
                .tools(McpToolCallbackProxy.proxyAll(mcpToolCallbacks))
                // 传递工具上下文参数，保证文件处理的隔离性（取消注释）
                .toolContext(Map.of("chatId", chatId))
                .call()
                .chatResponse();
        String result = chatResponse.getResult().getOutput().getText();
        // 记录响应结果日志
        log.info("TravelApp.chat({}): {}", chatId, result);
        return result;
    }

    /**
     * 与AI大模型流式对话接口
     * @param chatId 对话ID
     * @param message 消息内容
     * @return 大模型响应Flux对象
     */
    public Flux<String> streamChat(String chatId, String message) {
        // 从toolCallbackProvider获得MCP服务产生的ToolCallback数组
        FunctionCallback[] mcpToolCallbacks = toolCallbackProvider.getToolCallbacks();
        // 返回响应式Flux对象
        Flux<String> result = chatClient.prompt()
                .user(message)
                .advisors(spec -> spec
                        .param(CHAT_MEMORY_CONVERSATION_ID_KEY, chatId)
                        .param(CHAT_MEMORY_RETRIEVE_SIZE_KEY, 10)
                )
                // 注册可以使用的工具
                .tools(aiTools)
                // 注册代理后的MCP服务工具
                .tools(McpToolCallbackProxy.proxyAll(mcpToolCallbacks))
                // 传递工具上下文参数，保证文件处理的隔离性（取消注释）
                .toolContext(Map.of("chatId", chatId))
                .stream() // 用stream代替call方法，表示以流式模式返回结果
                .content();
        // 记录响应结果日志
        log.info("TravelApp.streamChat({}): {}", chatId, result);
        return result;
    }
}
